#ifndef _CUDA__HELPERS_CUH__
#define _CUDA__HELPERS_CUH__

#include <cuda.h>
#include "config.hxx"

template <typename scalar_t> 
void __global__ _gcn_normalize_csr_kernel(
  const int * __restrict__ row_offset,
  const int * __restrict__ adj_list,
  const int * __restrict__ odegree,
  const scalar_t * __restrict__ weight_list,
  scalar_t * __restrict__ output,
  const int nvertex )
{
  const int Tid = threadIdx.x + blockIdx.x * blockDim.x;
  const int Stride = blockDim.x * gridDim.x;
  for ( int v=Tid; v < nvertex; v += Stride )
  {
    for ( int indptr = row_offset[v]; indptr < row_offset[v+1]; ++ indptr )
    {
      int u = adj_list[indptr];
      // if (u == v) continue;
      auto uv = weight_list[indptr];
      int v_deg = odegree[v];
      int u_deg = odegree[u];
      // printf("uv[%d] : d[%d] = %d, d[%d] = %d\n", indptr, v, v_deg, u, u_deg);
      output[indptr] = uv / sqrtf(v_deg*u_deg);
    }
  }
}


template <typename scalar_t> 
void __global__ _gcn_normalize_coo_kernel(
  const int * __restrict__ src_list,
  const int * __restrict__ dst_list,
  const int * __restrict__ odegree,
  const scalar_t * __restrict__ weight_list,
  scalar_t * __restrict__ output,
  const int nedge )
{
  const int Tx = threadIdx.x + blockDim.x*blockIdx.x;
  const int Stride = blockDim.x * gridDim.x;

  for ( int i=Tx; i < nedge; i += Stride )
  {
    auto src = src_list[i];
    auto dst = dst_list[i];
    auto ew  = weight_list[i];
    
    auto sdeg = odegree[src];
    auto ddeg = odegree[dst];

    output[i] = ew / sqrtf(sdeg*ddeg);
  }
}

void __global__ _degree_gather(
  const int * __restrict__ vtx_list,
  int * __restrict__ degrees,
  const int nvertex, const int nedge)
{
  const int Tx = threadIdx.x;
  const int Bx = blockIdx.x;
  const int Stride = blockDim.x;

  const int SmSize = 2048; 
  __shared__ int sm[SmSize];

  for ( int i=Tx; i<SmSize; i+=Stride ) sm[i] = 0;

  for ( int base=Bx*SmSize; base < nvertex; base+=SmSize*gridDim.x)
  {
    for ( int i=Tx; i<nedge; i+=Stride )
    {
      int v = vtx_list[i];
      if (v>=base && v<base+SmSize) atomicAdd(sm+v-base, 1);
    }
    __syncthreads();
    for ( int i=Tx; i<SmSize; i+=Stride )
    {
      atomicAdd(degrees+base+i, sm[i]);
    }
    __syncthreads();
  }
}



#endif